#pragma once
#include "./custom_type.h"
#include "./data_transfer.h"
#include "./kernel_const.h"
#include "./kernel_utils.h"
#include "kernel_operator.h"


template<
    uint32_t L1_M0,
    uint32_t L1_N0
>
[aicore] inline __attribute__((always_inline)) void BatchMatrixPadding(
    layoutType layout,  
    int64_t zeroPaddingM, 
    int64_t zeroPaddingN,
    int64_t batchCount,
    __gm__ int64_t*  d_validM, 
    __gm__ int64_t*  d_validN, 
    __gm__ half**  d_matrixPointer, 
    // padding 相关
    __gm__ uint8_t *d_is_padding, 
    __gm__ half** d_matrixPointerPadding, 
    __gm__ int64_t *d_valid_padding_lda, 
    uint8_t paddingDir  // layout0: 0-Zz, 1-Nz; layout1: 0-Nn, 1-Zn
){

    // 申请UB完整空间
    AscendC::TBuf<AscendC::TPosition::VECIN> ub_tensor;
    AscendC::TPipe ub_pipe;
    ub_pipe.InitBuffer(ub_tensor, UB_BYTES);
    AscendC::LocalTensor<uint8_t> ub_tensor = ub_tensor.Get<uint8_t>();
    ub_pipe.Destroy();
    
    // 双缓冲开辟
    static constexpr uint32_t ub_paddingPingpongNum = 2;
    AscendC::LocalTensor<half> ub_paddingBuf[ub_paddingPingpongNum];

    // 划分缓冲区
    uint64_t ub_paddingPingpongBytes = UB_BYTES / ub_paddingPingpongNum;
    // 缓冲区的half元素个数
    uint64_t ub_paddingPingpongSize = ub_paddingPingpongBytes / sizeof(half);
    // 取偏移
    #pragma unroll
    for(uint32_t i = 0; i < ub_paddingPingpongNum; i++){
        ub_paddingBuf[i] = ub_tensor[ i * ub_paddingPingpongBytes ].template ReinterpretCast<half>();
    }
    
    // 缓冲区首次set
    #pragma unroll
    for(uint32_t i = 0; i < ub_paddingPingpongNum; i++){
        AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>((event_t)(i));
    }

    // 全局内存中需要padding的矩阵
    AscendC::GlobalTensor<half> gm_matrix;
    // padding后的矩阵地址
    AscendC::GlobalTensor<half> gm_matrixPadding;

    // 不同矩阵统一的padding参数
    // block是对于当前矩阵，每次要搬运到L1的数据大小，block中连续的维度这里为blockRowSize，另一维度为blockRowNum
    // 分块行长度（RowSize是连续存储的方向）
    int64_t blockRowSize;
    // 分块行数
    int64_t blockRowNum;
    // 根据M*N大小的矩阵行列优先来决定
    if(layout == RowMajor){
        blockRowSize = L1_N0; 
        blockRowNum = L1_M0; 
    }else{
        blockRowSize = L1_M0; 
        blockRowNum = L1_N0; 
    }
    int64_t blockSize = blockRowNum * blockRowSize;

    // 不同矩阵需要切换到padding参数
    // 对于同一逻辑矩阵，存在行优先和列优先存储的情况，连续存储的那个维度为matrixRowSize，另一维度为matrixRowNum
    // 原矩阵两行首间隔
    int64_t gm_stride = 0; 
    // 原矩阵同一行数据在 padding 后每段 blockRow 之间的的stride（根据Zz/Nz的时候来切换）
    int64_t gm_padding_stride = 0;
    // 连续存储方向上有效数据多长
    int64_t matrixRowSize = 0;
    // 有多少个连续存储方向的行
    int64_t matrixRowNum = 0;
    // 连续存储方向上能分多少块
    int64_t matrixRowSizeBlockLoops = 0; 
    // 列方向上分多少块
    int64_t matrixRowNumBlockLoops = 0; 
    // 当前ub缓冲区是否能转下矩阵的一行
    uint8_t isMatrixRowExceedBuf = 0;
    // 连续存储行的方向上能分多少任务（isMarixRowExceedBuf == 1的时候才需要用）
    int64_t matrixRowSizeTaskLoops = 0;
    // 多行方向上能分多少任务（ismatrixRowExceedBuf == 0或1 都需要用）
    int64_t matrixRowNumTaskLoops = 0;
    // 每个padding任务对应一次缓冲区使用
    // 当isMatrixRowExceedBuf == 1，计算一行需要分多少次来padding完
    int64_t tasksNumPerRow = 0;
    // 当isMatrixRowExceedBuf == 0，计算一个缓冲区能装多少矩阵行
    int64_t rowNumPerBuf = 0;

    // 当前的aiv号 返回0-39
    uint32_t aivIdx = AscendC::GetBlockIdx();
    // aiv总数
    uint32_t aivSum = AscendC::GetBlockNum() * AscendC::GetSubBlockNum();
    
    // 分核计算需要的参数
    // 当前aicore计算的第几个任务（0号开始）
    uint32_t cur_aicore_task = 0;
    // 从第0个矩阵开始，目前遍历了多少任务（1开始），用具计算当前遍历的任务是当前矩阵的第几个任务
    uint32_t cur_tasks_sum = 0;
    // 当前正在计算的矩阵号（0开始）
    uint32_t batchNum = -1;
    // 循环上限，不断更新
    int64_t loopSum = 0;

    // 分核计算
    for(uint32_t loopIdx = 0; loopIdx < loopSum+1; loopIdx++){
        
        // 矩阵切换
        if( loopIdx == loopSum ){ // 此时遍历到下一矩阵的0号任务
            batchNum++;
            if(batchNum < batchCount){
                if(d_is_padding[batchNum]){
                    cur_tasks_sum = loopSum;
                    gm_matrix.SetGlobalBuffer( (__gm__ half*)d_matrixPointer[batchNum] );
                    gm_matrixPadding.SetGlobalBuffer( (__gm__ half*)d_matrixPointerPadding[batchNum] );
                    gm_stride = (layout == RowMajor ? zeroPaddingN : zeroPaddingM);
                    matrixRowSize = (layout == RowMajor ? d_validN[batchNum] : d_validM[batchNum] );
                    matrixRowNum = (layout == RowMajor ? d_validM[batchNum] : d_validN[batchNum] );
                    matrixRowSizeBlockLoops = CeilDiv<int64_t>(matrixRowSize, blockRowSize); 
                    matrixRowNumBlockLoops = CeilDiv<int64_t>(matrixRowNum, blockRowNum);
                    if(paddingDir == 0){
                        gm_padding_stride = blockSize; 
                    }else{
                        gm_padding_stride = matrixRowNumBlockLoops * blockSize;
                    }
                    isMatrixRowExceedBuf = (matrixRowSize * sizeof(half) > ub_paddingPingpongBytes);
                    // 计算当前矩阵的padding分成多少个任务，每个任务分配给一个AIV，对应一个缓冲区
                    if(isMatrixRowExceedBuf){ // 当一个优先存储的行大于ub缓冲区，每次任务只能padding一行的一部分
                        tasksNumPerRow = CeilDiv<int64_t>(matrixRowSize * sizeof(half), ub_paddingPingpongBytes);
                        matrixRowSizeTaskLoops = tasksNumPerRow;
                        matrixRowNumTaskLoops = matrixRowNum;
                        loopSum += matrixRowSizeTaskLoops * matrixRowNumTaskLoops;
                    }else{ // 否则，每次任务至少可以padding一行，需要计算缓冲区能装下多少矩阵行的行
                        // 注意这里ub省出一个块的行长的空间，这是为了在ub->gm的时候，最后一行的最后一个块的行长搬运的时候让ub不要越界
                        rowNumPerBuf = (ub_paddingPingpongBytes - sizeof(half) * blockRowSize ) / (matrixRowSize * sizeof(half));
                        matrixRowNumTaskLoops = CeilDiv<int64_t>(matrixRowNum, rowNumPerBuf);
                        loopsum += matrixRowNumTaskLoops;
                    }
                }else{
                    // 此时当前循环没有对应一个任务块
                    // 还要判断下一个矩阵而不是结束循环
                    loopIdx--;
                    continue;
                }
            }else{ // 最后一次
                continue;
            }
        }

        // 将当前padding task分给一个aiv处理
        if(loopIdx % aivSum !=aivIdx ){
            continue;
        }

        // 当前任务是当前矩阵的第几个padding任务，根据isMatrixRowExceedBuf定位到原矩阵某一位置
        uint64_t taskIdx = loopIdx - cur_tasks_sum;
        // 当前任务对应哪个缓冲区
        uint8_t cur_task_pp_idx = cur_aicore_task % ub_paddingPingpongNum;

        // 当一个优先存储的行大于ub一半，每次任务只能padding一行的一部分
        if(isMatrixRowExceedBuf){

            // 当前任务的原矩阵横纵坐标号
            uint64_t matrixRowSizeTaskIdx = taskIdx % matrixRowSizeTaskLoops;
            uint64_t matrixRowNumTaskIdx = taskIdx / matrixRowSizeTaskLoops;
            // 当前任务的原矩阵横纵起始地址
            uint64_t matrixRowSizeTaskAddr = matrixRowSizeTaskIdx * ub_paddingPingpongSize;
            uint64_t matrixRowNumTaskAddr = matrixRowNumTaskIdx;
            uint64_t matrixTaskAddr = matrixRowNumTaskAddr * gm_stride + matrixRowSizeTaskAddr; 
            // 当前需要搬运的实际数据大小
            uint64_t gm_matrixRowSizeActual = (matrixRowSizeTaskIdx == matrixRowSizeTaskLoops-1 ? 
                                                matrixRowSize - matrixRowSizeTaskAddr
                                                : ub_paddingPingpongSize);
            uint64_t gm_matrixRowNumActual = 1;

            // 当前任务的第0段blockRow所在的块坐标
            uint64_t matrixRowSizeBlockIdx = matrixRowSizeTaskAddr / blockRowSize; 
            uint64_t matrixRowNumBlockIdx = matrixRowNumTaskAddr / blockRowNum;
            // 当前任务的第0段blockRow在块内的行号
            uint64_t matrixRowNumBlockInnerIdx = matrixRowNumTaskAddr % blockRowNum;
            // 当前任务的第0段blockRow的目的地址
            uint64_t matrixTaskPaddingAddr; 
            if(paddingDir == 0){ // Zz/Nn
                matrixTaskPaddingAddr = matrixRowNumBlockIdx * matrixRowSizeBlockLoops * blockSize
                                      + matrixRowSizeBlockIdx * blockSize
                                      + matrixRowNumBlockInnerIdx * blockRowSize;
            }else{ // Nz/Zn
                matrixTaskPaddingAddr = matrixRowSizeBlockIdx * matrixRowNumBlockLoops *  blockSize
                                      + matrixRowNumBlockIdx * blockSize
                                      + matrixRowNumBlockInnerIdx * blockRowSize;
            }

            // 循环间流水
            AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>((event_t)(cur_task_pp_idx));

            // 数据搬运：gm->ub
            Gm2Ub(
                ub_paddingBuf[cur_task_pp_idx], 
                gm_matrix[matrixTaskAddr], 
                gm_matrixRowNumActual, 
                gm_matrixRowSizeActual, 
                gm_stride
            );

            // 循环内流水
            AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE3>((event_t)(cur_task_pp_idx));
            AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE3>((event_t)(cur_task_pp_idx));

            // 数据搬运：ub->gm，分段搬运，每段长度为blockRowSize，最终分块间格式为Zz/Nz
            Ub2Gm(
                gm_matrixPadding[matrixTaskPaddingAddr], 
                ub_paddingBuf[cur_task_pp_idx], 
                CeilDiv<uint64_t>(gm_matrixRowSizeActual, blockRowSize), 
                blockRowSize, 
                gm_padding_stride
            );
            
            // 循环间流水
            AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>((event_t)(cur_task_pp_idx));
        
        }else{ // 否则，每次任务至少可以padding一行，缓冲区能装下几个完整几行，每次任务就padding几行

            // 将当前矩阵的padding任务号定位到某一行块的起始位置，搬运至少一行
            // 当前任务的原矩阵任务号
            uint64_t matrixRowNumTaskIdx = taskIdx;
            // 当前任务的原矩阵的纵向起始地址
            uint64_t matrixRowNumTaskAddr = matrixRowNumTaskIdx * rowNumPerBuf;
            uint64_t matrixTaskAddr = matrixRowNumTaskAddr * gm_stride;
            // 当前需要搬运的实际数据大小
            uint64_t gm_matrixRowSizeActual = matrixRowSize;
            uint64_t gm_matrixRowNumActual = (matrixRowNumTaskIdx == matrixRowNumTaskLoops-1 ? 
                                                matrixRowNum - matrixRowNumTaskAddr
                                                : rowNumPerBuf );
            
            // 当前任务中各行的第0段blockRow所在的块坐标 
            uint64_t matrixRowSizeBlockIdx; 
            uint64_t matrixRowNumBlockIdx; 
            // 当前任务中各行第0段blockRow在块内的行号
            uint64_t matrixRowNumBlockInnerIdx; 
            // 当前任务中各行的第0段blockRow的目的地址
            uint64_t matrixTaskPaddingAddr; 

            // 循环间流水
            AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>((event_t)(cur_task_pp_idx));

            // 数据搬运：gm->ub
            Gm2Ub(
                ub_paddingBuf[cur_task_pp_idx], 
                gm_matrix[matrixTaskAddr], 
                gm_matrixRowNumActual, 
                gm_matrixRowSizeActual, 
                gm_srcStride
            );

            // 循环内流水
            AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE3>((event_t)(cur_task_pp_idx));
            AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE3>((event_t)(cur_task_pp_idx));

            // 数据搬运：ub->gm，
            // 循环搬运每一行
            // 每行分段搬运，每段长度为blockRowSize，最终分块间格式为Zz/Nz
            for(uint32_t rowIdx = 0; rowIdx < gm_matrixRowNumActual; rowIdx++){
                matrixRowSizeBlockIdx = 0;
                matrixRowNumBlockIdx = (matrixRowNumTaskAddr + i) / blockRowNum;
                matrixRowNumBlockInnerIdx = (matrixRowNumTaskAddr + i) % blockRowNum;
                if(paddingDir == 0){ // Zz/Nn
                    matrixTaskPaddingAddr = matrixRowNumBlockIdx * matrixRowSizeBlockLoops * blockSize
                                          + matrixRowSizeBlockIdx * blockSize
                                          + matrixRowNumBlockInnerIdx * blockRowSize;
                }else{ // Nz/Zn
                    matrixTaskPaddingAddr = matrixRowSizeBlockIdx * matrixRowNumBlockLoops *  blockSize
                                          + matrixRowNumBlockIdx * blockSize
                                          + matrixRowNumBlockInnerIdx * blockRowSize;
                }
                Ub2Gm(
                    gm_matrixPadding[matrixTaskPaddingAddr], 
                    ub_paddingBuf[cur_task_pp_idx][i * matrixRowSize], 
                    CeilDiv<uint64_t>(gm_matrixRowSizeActual, blockRowSize), 
                    blockRowSize, 
                    gm_padding_stride
                );
            }

            // 循环间流水
            AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>((event_t)(cur_task_pp_idx));
            
        }

        cur_aicore_task ++;

    }

    // 缓冲区最后一次wait
    #pragma unroll
    for(uint32_t i = 0; i < ub_padingPingpongNum; i++){
        AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>((event_t)(i));
    }
}